from pathlib import Path
from typing import Union, Tuple, Optional, Sequence

import PIL
import mmengine
import numpy as np
from PIL import Image
from ultralytics import YOLO
from mmdeploy.apis import inference_model

default_classes = [
    'intact',
    'slight',
    'severe',
    'collapse']


def PIDNet(img: Union[str, np.ndarray],
           model_cfg: Union[str, mmengine.Config] = 'weights/PIDNet/pidnet-s_2xb6-120k_512x512-roadcracks.py',
           deploy_cfg: Union[str, mmengine.Config] = 'weights/PIDNet/segmentation_onnxruntime_static-512x512.py',
           backend_files=None,
           device: str = 'cpu') -> list:
    """
    PIDNet
    :param model_cfg: 模型配置文件
    :param deploy_cfg: 模型导出时的配置文件
    :param backend_files: 模型文件
    :param img: 图片 Input image file or numpy array for inference or url/image
    :param device: 设备
    :return: 图像数组
    """
    if backend_files is None:
        backend_files = ['weights/PIDNet/pidnet-s_2xb6-120k_512x512-roadcracks.onnx']
    results = inference_model(model_cfg, deploy_cfg, backend_files, img, device)
    images = []
    for result in results:
        image = result.get('pred_sem_seg').data.squeeze().numpy().astype(np.uint8)
        images.append(image)
    return images


def yolov8_seg(source: PIL.Image, imgsz=None, classes=None,
               conf=0.25, iou=0.5,
               model_path: Union[str, Path] = "weights/YOLO-Seg/yolo-v8-seg-n.onnx") -> Tuple[
    bool, list, list, Optional[list], Optional[list]]:
    """

    :param source: 图片源
    :param imgsz: 图像大小
    :param classes: 类别
    :param conf: 置信度
    :param iou: iou
    :param model_path: 模型位置
    :return: isExist, labels, scores, boxes, masks
    """
    # Load  model
    # if imgsz is None:
    imgsz = [1024, 1024]
    if classes is None:
        classes = default_classes
    model = YOLO(model_path, task='segment')
    # Run inference on the source
    results = model(source, device='cpu', imgsz=imgsz, conf=conf, iou=iou,
                    classes=[default_classes.index(x) for x in classes if
                             x in default_classes])  # list of Results objects
    result = results[0]
    labels = result.boxes.cls.cpu().numpy().astype(np.uint8).tolist()
    scores = result.boxes.conf.cpu().numpy().tolist()
    isExist = True
    boxes = None
    masks = None
    if result.masks is None:
        isExist = False
    else:
        boxes = result.boxes.xyxy.cpu().numpy().tolist()
        masks = result.masks.xy
    return isExist, labels, scores, boxes, masks


if __name__ == '__main__':
    classes = [
        'slight',
        'severe',
        'collapse']
    re = yolov8_seg(Image.open(r"D:\Users\13055\Desktop\IMG_20240627_152843.jpg"), classes=classes)
